package com.iflytek.aiui.websocketdemo;

import java.io.RandomAccessFile;
import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.Base64;
import java.util.Date;
import java.util.Locale;
import java.util.Objects;
import java.util.TimeZone;
import java.util.concurrent.CountDownLatch;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import org.java_websocket.WebSocket.READYSTATE;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.drafts.Draft;
import org.java_websocket.handshake.ServerHandshake;

import com.alibaba.fastjson.JSONObject;

/**
 * 运行方法：运行main()，控制台输出结果
 * 
 * @author Administrator
 *
 */
public class AIUIV2WsDemo {
    // 服务地址
    private static final String BASE_URL = "wss://aiui.xf-yun.com/v2/aiint/ws";

    // 应用配置
    private static final String APPID = "";        
    private static final String APIKEY = "";
    private static final String APISECRET= "";

    // 场景
    private static final String SCENE = "main_box";

    /**
     * 用于设置发送文本请求还是音频请求
     * 请求类型: text/audio
     * 
     * 音频请求需要先设置FILE_PATH
     * 当前音频格式默认pcm，修改音频格式需要修改audioReq中的payload中音频相关参数
     */
    private static final String REQ_TYPE = "audio";

    // 请求文本
    private static String text = "你是谁";
    // 文本请求参数
    private static String textReq = "{\"header\":{\"sn\":\"4b37ccbb6679c34389176ed5bc920196\",\"appid\":\"" + APPID + "\",\"stmid\":\"text-stream-1\",\"status\":3,\"scene\":\"" + SCENE + "\"},\"parameter\":{\"nlp\":{\"nlp\":{\"compress\":\"raw\",\"format\":\"json\",\"encoding\":\"utf8\"},\"sub_scene\":\"cbm_v45\",\"new_session\":true}},\"payload\":{\"text\":{\"compress\":\"raw\",\"format\":\"plain\",\"text\":\""+Base64.getEncoder().encodeToString(text.getBytes())+"\",\"encoding\":\"utf8\",\"status\":3}}}";

    // 每帧音频数据大小，单位字节
    private static final int CHUNCKED_SIZE = 1024;
    // 音频文件位置
    private static final String FILE_PATH = "D:/temp/audio/date.pcm";
    private static String audioReq = "{\"header\":{\"sn\":\"4b37ccbb6679c34389176ed5bc920196\",\"appid\":\"" + APPID + "\",\"stmid\":\"test-stream-1\",\"status\":3,\"scene\":\"" + SCENE + "\"},\"parameter\":{\"nlp\":{\"nlp\":{\"compress\":\"raw\",\"format\":\"json\",\"encoding\":\"utf8\"},\"sub_scene\":\"cbm_v45\",\"new_session\":true}},\"payload\":{\"audio\":{\"encoding\":\"raw\",\"sample_rate\":1600,\"channels\":1,\"bit_depth\":16}}}";

    // main()方法，直接运行，控制台输出服务端结果
    public static void main(String[] args) throws Exception {
        if (!Objects.equals(REQ_TYPE, "text") && !Objects.equals(REQ_TYPE, "audio")) {
            throw new RuntimeException("REQ_TYPE非法");
        }

        URI url = new URI(BASE_URL + getHandShakeParams());
        DraftWithOrigin draft = new DraftWithOrigin(getWebsocketUrlOrigin(BASE_URL));
        CountDownLatch countDownLatch = new CountDownLatch(1);
        MyWebSocketClient client = new MyWebSocketClient(url, draft, countDownLatch);
        client.connect();

        while (!client.getReadyState().equals(READYSTATE.OPEN)) {
            Thread.sleep(50);
        }

        if (Objects.equals(REQ_TYPE, "text")) {
            // 文本请求
            client.send(textReq);
        } else {
            // 音频请求
            byte[] bytes = new byte[CHUNCKED_SIZE];
            try (RandomAccessFile raf = new RandomAccessFile(FILE_PATH, "r")) {
                int len = -1;
                int readIndex = 0;
                int fileLength = (int) raf.length();
                int lastIndex = fileLength / CHUNCKED_SIZE + (fileLength % CHUNCKED_SIZE == 0 ? 0 : 1);
                while ((len = raf.read(bytes)) != -1) {
                    readIndex += 1;
                    if (len < CHUNCKED_SIZE) {
                        bytes = Arrays.copyOfRange(bytes, 0, len);
                    }

                    JSONObject req = JSONObject.parseObject(audioReq);
                    JSONObject audio = req.getJSONObject("payload").getJSONObject("audio");
                    // 设置音频格式
                    audio.put("encoding", "raw");
                    audio.put("sample_rate", 16000);
                    audio.put("channels", 1);
                    audio.put("bit_depth", 16);
                    audio.put("frame_size", bytes.length);

                    // 首帧
                    int status = 0;
                    if (lastIndex == 1) {
                        // 一次传完
                        status = 3;
                    } else if (readIndex == lastIndex) {
                        // 结束帧
                        status = 2;
                    } else if (readIndex != 1) {
                        // 中间帧
                        status = 1;
                    }
                    req.getJSONObject("header").put("status", status);
                    audio.put("status", status);
                    audio.put("audio", new String(Base64.getEncoder().encode(bytes)));

                    if (client.isClosed()) {
                        System.out.println("连接异常关闭");
                        return;
                    }
                    client.send(req.toJSONString());
                    Thread.sleep(40);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        // 等待交互完成
        countDownLatch.await();
        if (!client.isClosed()) {
            // 主动断开连接，如果不主动断开60秒无交互之后服务端会主动断开连接
            client.close(1000);
        }
    }

    // 拼接握手参数
    private static String getHandShakeParams() throws Exception {
        URI uri = URI.create(BASE_URL);
        String host = uri.getHost();
        String path = uri.getPath();
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        // 拼接
        String preStr = "host: " + host + "\n" + "date: " + date + "\n" + "GET " + path + " HTTP/1.1";
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(APISECRET.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"",
                APIKEY, "hmac-sha256", "host date request-line", sha);
        // 拼接鉴权需要的查询参数
        return String.format("?authorization=%s&date=%s&host=%s", URLEncoder
                .encode(Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)), "UTF-8"),
                URLEncoder.encode(date, "UTF-8"), URLEncoder.encode(host, "UTF-8"));
    }

    private static String getWebsocketUrlOrigin(String websocketUrl) {
        websocketUrl = websocketUrl.replaceAll("ws[s]?://", "");
        return "http://" + websocketUrl.substring(0, websocketUrl.indexOf("/"));
    }

    private static class MyWebSocketClient extends WebSocketClient {

        private CountDownLatch countDownLatch;

        public MyWebSocketClient(URI serverUri, Draft protocolDraft, CountDownLatch countDownLatch) {
            super(serverUri, protocolDraft);
            this.countDownLatch = countDownLatch;
        }

        @Override
        public void onOpen(ServerHandshake handshake) {
            // 会阻塞onMesage, 如果要在这里发送数据需要另起线程
            System.out.println("打开连接, code:" + handshake.getHttpStatusMessage());
        }

        @Override
        public void onMessage(String msg) {
            System.out.println("原始结果：" + msg);
            JSONObject response = JSONObject.parseObject(msg);
            // 结果解析
            JSONObject payload = response.getJSONObject("payload");
            if (payload != null) {
                // 语义结果
                JSONObject nlp = payload.getJSONObject("nlp");
                if (nlp != null) {
                    System.out.println("语义结果 nlp text：" + new String(Base64.getDecoder().decode(nlp.getString("text"))));
                }
                JSONObject iat = payload.getJSONObject("iat");
                if (iat != null) {
                    System.out.println("识别结果 iat text：" + new String(Base64.getDecoder().decode(iat.getString("text"))));
                }
            }

            JSONObject header = response.getJSONObject("header");
            if (!header.containsKey("status")) {
                return;
            }
            int status = header.getIntValue("status");
            if (status == 2) {
                // 接收到最后一帧结果
                countDownLatch.countDown();
            }
        }

        @Override
        public void onError(Exception e) {
            System.out.println("连接发生错误：" + e.getMessage() + ", " + new Date());
            e.printStackTrace();
        }

        @Override
        public void onClose(int code, String message, boolean b) {
            System.out.println("连接已关闭" + "," + new Date());
            countDownLatch.countDown();
        }

        @Override
        public void onMessage(ByteBuffer bytes) {
            try {
                String responseText = new String(bytes.array(), "utf-8");
                System.out.println("服务端返回：" + responseText);
            } catch (UnsupportedEncodingException e) {
                e.printStackTrace();
            }
        }
    }
}